import numpy as np

import astropy.units as u


def newton_generation(initial_value, e_1, M_1, tolerance):
    """
    牛顿迭代求解开普勒方程 E - e sin E = M
    :param initial_value: 开始迭代的值，可以为任意值，下面代码中取了上一个时刻的E为初始值
    :param e_1: 加了_1后缀是为了和下面代码中的e和M进行区分
    :param M_1: 同上
    :param tolerance: 迭代到误差小于什么时候停止，下面代码中取了1e-6
    :return: 计算结果
    """
    difference = 2 * tolerance
    old_value = initial_value
    new_value = initial_value
    while difference > tolerance:
        new_value = old_value - (old_value - e_1 * np.sin(old_value) * u.rad - M_1) / (1 - e_1 * np.cos(old_value))
        difference = abs(old_value - new_value)
        old_value = new_value

    return new_value


utc_time = '2022-03-16T11:06:00'
def get_orbit_elements(r_0, r_dot_0):

    # 下面计算轨道根数，先得到目标的初始时刻位置和速度
    object_initial_location = r_0 * u.m
    object_initial_speed = r_dot_0 * u.m / u.s

    # 由目标初始位置和速度得到h
    h_vector = np.outer(object_initial_location, object_initial_speed)
    h_scalar = np.linalg.norm(h_vector)

    # 由h矢量得到方向向量R
    R_vector = h_vector / h_scalar

    # 再由方向向量得到i和Ω
    i = np.arccos(R_vector[-1])
    omega = np.arcsin(R_vector[0] / np.sin(i))

    # 引入μ，参考书本附录，模型是WGS84, 因为测站的坐标转换中，第32行代码里from_geodetic，默认选择了WGS84模型
    mu = 398600.4418 * 1e9 * (u.m ** 3) / ((u.s) ** 2)

    # 计算目标初始时刻位置和速度的大小
    r_scalar = np.linalg.norm(object_initial_location)
    v_scalar = np.linalg.norm(object_initial_speed)

    # 计算轨道半长轴
    a = 1 / (2 / r_scalar - v_scalar ** 2 / mu)

    # 计算E, e, M
    E = np.arctan2(np.dot(object_initial_speed, object_initial_location) / np.sqrt(mu * a), (1 - r_scalar / a))
    e = (1 - r_scalar / a) / np.cos(E)
    M = E - e * np.sin(E) * u.rad

    # 计算向量P和Q，用于后续得到目标位置和速度的矢量
    P_vector = np.cos(E) / r_scalar * object_initial_location - np.sqrt(a / mu) * np.sin(E) * object_initial_speed
    Q_vector = (np.sin(E) / r_scalar * object_initial_location + np.sqrt(a / mu) * (np.cos(E) - e) * object_initial_speed)

    # 计算ω， 但后续没用到
    w = np.arctan2(P_vector[-1], Q_vector[-1])

    return np.array([a, e, i, omega, w, M])
